#!/usr/bin/env python3
"""
mapAnalyzer.py

解析 txt 格式的系统 maps 文件（每行包含形如
"1aa0000000-1aa00c0000 rw-p 00000000 00:00 0   [anon:ArkTS Heaphuge object space]"
的记录），
根据十六进制的起始地址和结束地址计算每个条目的大小，
以方括号内的内容作为内存区名，统计相同内存区的总大小并输出。

用法示例：
  python3 mapAnalyzer.py /path/to/maps.txt
  cat /path/to/maps.txt | python3 mapAnalyzer.py
"""

import argparse
import re
import sys
from collections import defaultdict


ADDR_RE = re.compile(r"^\s*([0-9a-fA-F]+)-([0-9a-fA-F]+)")
BRACKET_RE = re.compile(r"\[(.*?)\]")


def parse_line(line):
	"""解析一行，返回 (name, size) 或 None。

	name: 方括号内的内容（去掉 []），若不存在则返回 'anonymous'
	size: end - start（整数，字节）
	"""
	m = ADDR_RE.match(line)
	if not m:
		return None
	try:
		start = int(m.group(1), 16)
		end = int(m.group(2), 16)
	except ValueError:
		return None
	size = max(0, end - start)

	b = BRACKET_RE.search(line)
	if b:
		name = b.group(1).strip()
	else:
		# 如果没有方括号，尝试用最后一列作为名字，否则设为 anonymous
		parts = line.strip().split()
		last = parts[-1] if parts else ''
		if last.startswith('[') and last.endswith(']'):
			name = last[1:-1]
		else:
			name = 'anonymous'

	return name, size


def human_readable(n):
	"""把字节数转成人类可读字符串，保留两位小数"""
	if n < 1024:
		return f"{n} B"
	for unit in ['KB', 'MB', 'GB', 'TB']:
		n /= 1024.0
		if n < 1024:
			return f"{n:.2f} {unit}"
	return f"{n:.2f} PB"


def analyze_lines(lines, include_prefix=None, exclude_prefix=None):
	"""聚合 lines 中的区域大小，可选根据前缀包含/排除过滤。

	include_prefix: 如果给定，则只统计 name 以此前缀开始的条目
	exclude_prefix: 如果给定，则排除 name 以此前缀开始的条目
	"""
	agg = defaultdict(int)
	total = 0
	for ln in lines:
		parsed = parse_line(ln)
		if not parsed:
			continue
		name, size = parsed

		# 过滤逻辑：先判断 include，再判断 exclude
		if include_prefix and not name.startswith(include_prefix):
			continue
		if exclude_prefix and name.startswith(exclude_prefix):
			continue

		agg[name] += size
		total += size
	return agg, total


def main():
	parser = argparse.ArgumentParser(description='Analyze maps-like txt and aggregate region sizes')
	parser.add_argument('file', nargs='?', help='maps txt file (omit to read from stdin)')
	parser.add_argument('-n', '--top', type=int, default=0, help='只显示前 N 个最大的内存区（默认全部）')
	parser.add_argument('--no-sort', action='store_true', help='不按大小排序，按内存区名输出')
	parser.add_argument('--include-prefix', help='只统计名字以此前缀开头的内存区')
	parser.add_argument('--exclude-prefix', help='排除名字以此前缀开头的内存区')
	parser.add_argument('--arkts', action='store_true', help='便捷开关：等同于 --include-prefix "anon:ArkTS" --exclude-prefix "anon:ArkTS Code"')
	parser.add_argument('--native', action='store_true', help='便捷开关：只统计名字以 "anon:native" 开头的内存区')
	parser.add_argument('--kotlin', action='store_true', help='便捷开关：只统计名字以 "anon:kotlin_native" 开头的内存区')
	args = parser.parse_args()

	if args.file:
		try:
			with open(args.file, 'r', encoding='utf-8', errors='replace') as f:
				lines = f.readlines()
		except Exception as e:
			print(f'无法打开文件 {args.file}: {e}', file=sys.stderr)
			sys.exit(2)
	else:
		lines = sys.stdin.readlines()

	# 过滤选项
	include_prefix = None
	exclude_prefix = None
	if args.arkts:
		include_prefix = 'anon:ArkTS'
		exclude_prefix = 'anon:ArkTS Code'
	if args.native:
		include_prefix = 'anon:native'
	if args.kotlin:
		include_prefix = 'anon:kotlin_native'
	if args.include_prefix:
		include_prefix = args.include_prefix
	if args.exclude_prefix:
		exclude_prefix = args.exclude_prefix

	agg, total = analyze_lines(lines, include_prefix=include_prefix, exclude_prefix=exclude_prefix)

	items = list(agg.items())
	if not args.no_sort:
		items.sort(key=lambda x: x[1], reverse=True)
	else:
		items.sort(key=lambda x: x[0])

	if args.top > 0:
		items = items[:args.top]

	# 输出表头
	print(f"{'Region':60} {'Bytes':>15} {'Human':>15}")
	print('-' * 96)
	for name, size in items:
		print(f"{name:60} {size:15d} {human_readable(size):>15}")

	print('-' * 96)
	print(f"{'Total':60} {total:15d} {human_readable(total):>15}")


if __name__ == '__main__':
	main()

